import cv2
import numpy as np
import torch
import torch.utils.data as Data

from model import RAS
from utils import trans_im, trans_gt

BATCH_SIZE = 1
EPOCHS_NUM = 200
LOG_FILE = "train_log.txt"


def get_train_data(start_image_id, end_image_id):
    image_num = end_image_id - start_image_id + 1
    x = np.zeros([image_num, 3, 500, 500])
    y = np.zeros([image_num, 1, 500, 500])
    for i in range(image_num):
        img_id = start_image_id + i
        im_path = "data/train/images/{:04d}.jpg".format(img_id)
        gt_path = "data/train/ground_truth_mask/{:04d}.png".format(img_id)
        x[i, :, :, :] = trans_im(im_path)
        y[i, :, :, :] = trans_gt(gt_path)

    x = torch.FloatTensor(x)
    y = torch.FloatTensor(y)
    torch_dataset = Data.TensorDataset(x, y)
    loader = Data.DataLoader(
        dataset=torch_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4,
    )
    return loader


if __name__ == "__main__":
    ras = RAS()
    device = torch.device("cpu")
    ras.to(device)

    train_loader = get_train_data(1, 700)

    # 清空原有日志文件内容
    with open(LOG_FILE, 'w') as f:
        f.write("Training Log\n")

    for epoch in range(99, EPOCHS_NUM):
        with open(LOG_FILE, 'a') as log_file:
            log_file.write(f"\nEpoch {epoch}:\n")
            for step, (batch_x, batch_y) in enumerate(train_loader):
                loss = ras.train(batch_x.to(device), batch_y.to(device))
                log_line = f"Epoch: {epoch}, Step: {step}, Loss: {loss}\n"
                log_file.write(log_line)
                if step % 20 == 0:
                    print(log_line.strip())
                if step % 100 == 0:
                    torch.save(ras.state_dict(), 'data/model/params.pkl')
            torch.save(ras.state_dict(), f'data/model/epoch_{epoch}_params.pkl')
